import json
import time
import re
import os
from openai import OpenAI
from docx import Document

# ========================== 1. 配置区 ==========================
INPUT_DOCUMENT_PATH = r"C:/Users/cassi/Desktop/bisai/word-master/豫洛阳-兰郑长干线-CPY-0790-BFGDGS-ZZSYQFGS.docx"
OUTPUT_JSON_PATH = r"C:/Users/cassi/Desktop/bisai/word-master/cross_reference_check.json"
SAMPLES_PER_TYPE = 3  # 每个负样本类型生成的变体数量（可调整）

# ========================== 2. 功能函数区 ==========================
def get_answer(message):
    """调用模型生成负样本"""
    try:
        client = OpenAI(
            api_key="sk-ma0HIkN3beK0mtThtDzIfRLN5IzpL4agmB9RSbpHk53h7BF5",
            base_url="https://www.chataiapi.com/v1",
        )
        response = client.chat.completions.create(
            model="gemini-2.5-pro",
            messages=message,
            temperature=0.7,
        )
        content = response.choices[0].message.content or ""
        if content.strip():
            print("\n*******************模型生成内容如下********************\n")
            print(content)
        return content
    except Exception as e:
        print(f"API请求失败：{str(e)}")
        return ""


# ========================== 2. 解析 Word ==========================
def parse_docx_sections(doc_path):
    """解析 Word 文档，识别主要表格区段"""
    doc = Document(doc_path)
    sections = {"basic_info": "", "risk_table": "", "prevention": "", "full_text": ""}
    current_section = None

    for para in doc.paragraphs:
        text = para.text.strip()
        if not text:
            continue
        sections["full_text"] += text + "\n"

        if re.search(r'高后果区基本信息|基本信息表', text):
            current_section = "basic_info"
        elif re.search(r'高后果区风险评价结果表|风险评价结果表|风险评价', text):
            current_section = "risk_table"
        elif re.search(r'人防措施|防护措施', text):
            current_section = "prevention"

        if current_section:
            sections[current_section] += text + "\n"

    # 合并表格文本
    for table in doc.tables:
        table_text = "\n".join(["  ".join([cell.text.strip() for cell in row.cells]) for row in table.rows])
        if re.search(r'高后果区基本信息|基本信息表', table_text):
            sections["basic_info"] += "\n" + table_text
        elif re.search(r'高后果区风险评价结果表|风险评价结果表|风险评价', table_text):
            sections["risk_table"] += "\n" + table_text
        elif re.search(r'人防|防护', table_text):
            sections["prevention"] += "\n" + table_text

    return sections

# ========================== 3. 信息提取 ==========================
def extract_section_leaders(sections):
    """提取区段长信息（支持提取多个，包含职位名称）"""
    leader_pattern = re.compile(
        r'(高后果区总区长|高后果区分区长|专职区长|专职区（段）长|区段长|负责人|责任人|管理人|巡护人)[：:\s]*([^\s，,。\n（）()]{2,8})'
    )
    prevention_matches = re.findall(leader_pattern, sections.get("prevention", ""))
    basic_info_matches = re.findall(leader_pattern, sections.get("basic_info", ""))
    
    # 过滤掉不完整的数据
    prevention_leaders = []
    for title, name in prevention_matches:
        if name and len(name.strip()) >= 2 and not any(char in name for char in ['等', '度', 'm', ')', '）']):
            prevention_leaders.append(f"{title}：{name.strip()}")
    
    basic_info_leaders = []
    for title, name in basic_info_matches:
        if name and len(name.strip()) >= 2 and not any(char in name for char in ['等', '度', 'm', ')', '）']):
            basic_info_leaders.append(f"{title}：{name.strip()}")
    
    # 去重并保持顺序
    prevention_leaders = list(dict.fromkeys(prevention_leaders))
    basic_info_leaders = list(dict.fromkeys(basic_info_leaders))
    
    return {
        "from_prevention_measures": prevention_leaders,
        "from_basic_info_table": basic_info_leaders
    }

def extract_locations(sections):
    """提取位置信息（支持提取多个，包含位置类型）"""
    loc_pattern = re.compile(
        r'(位置|位于|行政区划|地址|地理位置|高后果区起点|高后果区终点)[：:\s]*([^\s，,。\n]{3,30})'
    )
    basic_info_matches = re.findall(loc_pattern, sections.get("basic_info", ""))
    risk_table_matches = re.findall(loc_pattern, sections.get("risk_table", ""))
    
    # 组合位置类型和具体位置
    basic_info_locations = [f"{loc_type}：{location}" for loc_type, location in basic_info_matches]
    risk_table_locations = [f"{loc_type}：{location}" for loc_type, location in risk_table_matches]
    
    # 去重并保持顺序
    basic_info_locations = list(dict.fromkeys(basic_info_locations))
    risk_table_locations = list(dict.fromkeys(risk_table_locations))
    
    return {
        "from_basic_info": basic_info_locations,
        "from_risk_table": risk_table_locations
    }

def extract_hca_numbers(sections):
    """提取高后果区编号（通用增强版）"""
    # 改进后的正则模式（兼容省份/线路/支线等多样写法）
    number_patterns = [
        # 通用完整格式，例如：
        r'[豫Y][\u4e00-\u9fa5A-Z0-9]+[-—_－]?[A-Z\u4e00-\u9fa5]*[-—_－]?(CPY[-_—]?\d{3,4}(?:[-_A-Z0-9]+)?)',

        # 简写或残缺格式（仅编号部分）
        r'(CPY[-_—]?\d{3,4}(?:[-_A-Z0-9]+)?)'
    ]

    found = set()
    nums = []
    text = sections.get("full_text", "")

    for p in number_patterns:
        for match in re.finditer(p, text):
            num = match.group(1) if match.lastindex else match.group(0)
            if num not in found:
                found.add(num)
                nums.append(num)

    return nums



# ========================== 5. 模型生成负样本 ==========================
def get_cross_reference_prompt(leaders, locations, hca_numbers, prompt_type, sample_idx):
    """生成交叉引用检查的提示词模板"""
    prompt_templates = {
        1: f"""任务：上下文内容一致性"区段长信息不一致"负样本（第{sample_idx}个变体）。
要求：
1. 修改区段长信息，使其在不同位置出现不一致（如人防措施中为"区段长：张三"，基本信息表中为"区段长：李四"）；
2. 保持其他信息不变（如位置、编号、时间等），仅让区段长信息产生矛盾；
3. 确保修改后的信息在各自位置内逻辑合理，仅让交叉引用时出现不一致；
4. 这是第{sample_idx}个变体，请生成与之前不同的具体姓名和表达方式。
原始信息：
---
人防措施：{', '.join(leaders["from_prevention_measures"]) if leaders["from_prevention_measures"] else "无数据"}
基本信息表：{', '.join(leaders["from_basic_info_table"]) if leaders["from_basic_info_table"] else "无数据"}
---
输出：仅返回修改后的不一致信息对比，格式为：
人防措施：[修改后的区段长信息]
基本信息表：[修改后的区段长信息]""",

        2: f"""任务：上下文内容一致性"位置信息不一致"负样本（第{sample_idx}个变体）。
要求：
1. 修改位置描述，使其在不同表格中出现不一致（如基本信息表为"位置：河南省洛阳市"，风险评价表为"位置：河南省郑州市"）；
2. 保持其他信息不变（如区段长、编号、时间等），仅让位置信息产生矛盾；
3. 确保修改后的位置描述在各自表格内逻辑合理，仅让交叉引用时出现不一致；
4. 这是第{sample_idx}个变体，请生成与之前不同的具体位置和表达方式。
原始信息：
---
基本信息表：{', '.join(locations["from_basic_info"]) if locations["from_basic_info"] else "无数据"}
风险评价表：{', '.join(locations["from_risk_table"]) if locations["from_risk_table"] else "无数据"}
---
输出：仅返回修改后的不一致信息对比，格式为：
基本信息表：[修改后的位置信息]
风险评价表：[修改后的位置信息]""",

        3: f"""任务：上下文内容一致性"高后果区编号不一致"负样本（第{sample_idx}个变体）。
要求：
1. 修改高后果区编号，使其在文档不同位置出现不一致（如开头为"CPY-0790"，后续为"CPY-0800"）；
2. 保持其他信息不变（如区段长、位置、时间等），仅让编号信息产生矛盾；
3. 确保修改后的编号在各自位置内格式正确，仅让交叉引用时出现不一致；
4. 这是第{sample_idx}个变体，请生成与之前不同的具体编号。
原始信息：
---
文档编号：{', '.join(hca_numbers) if hca_numbers else "无数据"}
---
输出：仅返回修改后的不一致信息对比，格式为：
文档开头编号：[修改后的编号]
文档后续编号：[修改后的编号]"""
    }
    return prompt_templates.get(prompt_type, "")

def generate_structured_samples_with_model(leaders, locations, hca_numbers, samples_per_type=3):
    """使用模型生成训练样本，每个类型生成多个变体"""
    samples = []
    
    # 检查是否有足够的数据生成样本
    has_leader_data = leaders["from_prevention_measures"] or leaders["from_basic_info_table"]
    has_location_data = locations["from_basic_info"] or locations["from_risk_table"]
    has_hca_data = len(hca_numbers) > 0
    
    if not (has_leader_data or has_location_data or has_hca_data):
        print("⚠️ 警告：未找到足够的数据生成交叉引用检查样本")
        return samples
    
    # 生成区段长不一致样本（多个变体）
    if has_leader_data:
        print(f"🔄 正在生成区段长信息不一致负样本（{samples_per_type}个变体）...")
        for i in range(samples_per_type):
            print(f"  → 生成第{i+1}/{samples_per_type}个区段长样本...")
            prompt = get_cross_reference_prompt(leaders, locations, hca_numbers, 1, i+1)
            chat_message = [{"role": "user", "content": prompt}]
            response = get_answer(chat_message)
            
            if response.strip():
                try:
                    # 解析模型返回的不一致信息
                    lines = response.strip().split('\n')
                    prevention_info = ""
                    basic_info = ""
                    
                    for line in lines:
                        if "人防措施：" in line:
                            prevention_info = line.replace("人防措施：", "").strip()
                        elif "基本信息表：" in line:
                            basic_info = line.replace("基本信息表：", "").strip()
                    
                    if prevention_info and basic_info:
                        # 构造用户问题
                        user_content = f"人防措施中涉及的区段长信息与高后果区基本信息表是否一致。\n\n人防措施：{prevention_info}\n高后果区基本信息：{basic_info}"
                        
                        # 构造助手回答
                        # 统一使用第一个找到的区段长信息作为正确描述
                        all_leaders = leaders["from_prevention_measures"] + leaders["from_basic_info_table"]
                        correct_leader = all_leaders[0] if all_leaders else "无数据"
                        assistant_content = f"检测到上下文内容一致性问题：区段长信息不一致。\n正确的描述应该是：{correct_leader}"
                        
                        samples.append({
                            "messages": [
                                {"role": "user", "content": user_content},
                                {"role": "assistant", "content": assistant_content}
                            ],
                            "images": []
                        })
                        print(f"    ✅ 区段长样本{i+1}生成成功")
                    else:
                        raise ValueError("无法解析模型返回的不一致信息")
                        
                except Exception as e:
                    print(f"    ⚠️ 区段长样本{i+1}处理异常：{e}")
            else:
                print(f"    ⚠️ 区段长样本{i+1}生成失败")
            
            time.sleep(1)  # API限流控制
    
    # 生成位置信息不一致样本（多个变体）
    if has_location_data:
        print(f"🔄 正在生成位置信息不一致负样本（{samples_per_type}个变体）...")
        for i in range(samples_per_type):
            print(f"  → 生成第{i+1}/{samples_per_type}个位置样本...")
            prompt = get_cross_reference_prompt(leaders, locations, hca_numbers, 2, i+1)
            chat_message = [{"role": "user", "content": prompt}]
            response = get_answer(chat_message)
            
            if response.strip():
                try:
                    # 解析模型返回的不一致信息
                    lines = response.strip().split('\n')
                    basic_info = ""
                    risk_info = ""
                    
                    for line in lines:
                        if "基本信息表：" in line:
                            basic_info = line.replace("基本信息表：", "").strip()
                        elif "风险评价表：" in line:
                            risk_info = line.replace("风险评价表：", "").strip()
                    
                    if basic_info and risk_info:
                        # 构造用户问题
                        user_content = f"位置信息与高后果区风险评价结果表是否一致。\n\n基本信息表：{basic_info}\n风险评价结果表：{risk_info}"
                        
                        # 构造助手回答
                        original_basic = ', '.join(locations["from_basic_info"]) if locations["from_basic_info"] else "无数据"
                        original_risk = ', '.join(locations["from_risk_table"]) if locations["from_risk_table"] else "无数据"
                        assistant_content = f"检测到上下文内容一致性问题：位置信息不一致。\n正确的描述应该是：\n基本信息表：{original_basic}\n风险评价结果表：{original_risk}"
                        
                        samples.append({
                            "messages": [
                                {"role": "user", "content": user_content},
                                {"role": "assistant", "content": assistant_content}
                            ],
                            "images": []
                        })
                        print(f"    ✅ 位置样本{i+1}生成成功")
                    else:
                        raise ValueError("无法解析模型返回的不一致信息")
                        
                except Exception as e:
                    print(f"    ⚠️ 位置样本{i+1}处理异常：{e}")
            else:
                print(f"    ⚠️ 位置样本{i+1}生成失败")
            
            time.sleep(1)  # API限流控制
    
    # 生成高后果区编号不一致样本（多个变体）
    if has_hca_data:
        print(f"🔄 正在生成高后果区编号不一致负样本（{samples_per_type}个变体）...")
        for i in range(samples_per_type):
            print(f"  → 生成第{i+1}/{samples_per_type}个编号样本...")
            prompt = get_cross_reference_prompt(leaders, locations, hca_numbers, 3, i+1)
            chat_message = [{"role": "user", "content": prompt}]
            response = get_answer(chat_message)
            
            if response.strip():
                try:
                    # 解析模型返回的不一致信息
                    lines = response.strip().split('\n')
                    start_num = ""
                    end_num = ""
                    
                    for line in lines:
                        if "文档开头编号：" in line:
                            start_num = line.replace("文档开头编号：", "").strip()
                        elif "文档后续编号：" in line:
                            end_num = line.replace("文档后续编号：", "").strip()
                    
                    if start_num and end_num:
                        # 构造用户问题
                        user_content = f"高后果区编号前后是否一致。\n\n文档开头编号：{start_num}\n文档后续编号：{end_num}"
                        
                        # 构造助手回答
                        # 即使提取到多个编号，也统一使用第一个作为正确编号
                        correct_num = hca_numbers[0] if hca_numbers else "无数据"
                        assistant_content = f"检测到上下文内容一致性问题：高后果区编号不一致。\n正确的描述应该是：\n文档各处编号：{correct_num}"
                        
                        samples.append({
                            "messages": [
                                {"role": "user", "content": user_content},
                                {"role": "assistant", "content": assistant_content}
                            ],
                            "images": []
                        })
                        print(f"    ✅ 编号样本{i+1}生成成功")
                    else:
                        raise ValueError("无法解析模型返回的不一致信息")
                        
                except Exception as e:
                    print(f"    ⚠️ 编号样本{i+1}处理异常：{e}")
            else:
                print(f"    ⚠️ 编号样本{i+1}生成失败")
            
            time.sleep(1)  # API限流控制
    
    return samples

# ========================== 6. 主程序 ==========================
if __name__ == "__main__":
    fname = os.path.basename(INPUT_DOCUMENT_PATH)
    print(f"\n{'='*80}\n正在处理文件: {fname}\n{'='*80}")

    # 解析文档
    sections = parse_docx_sections(INPUT_DOCUMENT_PATH)
    leaders = extract_section_leaders(sections)
    locations = extract_locations(sections)
    hca_numbers = extract_hca_numbers(sections)
    
    # 显示提取的信息
    print(f"\n📊 提取的信息统计：")
    print(f"   - 区段长信息：人防措施 {len(leaders['from_prevention_measures'])} 条，基本信息表 {len(leaders['from_basic_info_table'])} 条")
    print(f"   - 位置信息：基本信息表 {len(locations['from_basic_info'])} 条，风险评价表 {len(locations['from_risk_table'])} 条")
    print(f"   - 高后果区编号：{len(hca_numbers)} 条")
    
    # 检查数据质量
    has_leader_data = leaders["from_prevention_measures"] or leaders["from_basic_info_table"]
    has_location_data = locations["from_basic_info"] or locations["from_risk_table"]
    has_hca_data = len(hca_numbers) > 0
    
    if not (has_leader_data or has_location_data or has_hca_data):
        print(f"\n⚠️ 警告：未找到足够的数据生成交叉引用检查样本")
        print("请检查文档是否包含区段长、位置或编号信息")
        exit()
    
    # 使用模型生成负样本
    print(f"\n🤖 开始使用模型生成交叉引用检查负样本...")
    print(f"📝 配置：每个负样本类型将生成 {SAMPLES_PER_TYPE} 个变体")
    samples = generate_structured_samples_with_model(leaders, locations, hca_numbers, SAMPLES_PER_TYPE)

    # 保存结果
    if samples:
        with open(OUTPUT_JSON_PATH, "w", encoding="utf-8") as f:
            json.dump(samples, f, ensure_ascii=False, indent=4)
        print(f"\n{'='*80}")
        print(f"✅ 处理完成！数据保存至：{OUTPUT_JSON_PATH}")
        print(f"共生成 {len(samples)} 个对话样本（每个包含负样本user + 正样本assistant）")
        
        # 统计各类型样本数量
        leader_count = sum(1 for s in samples if "区段长信息不一致" in s["messages"][1]["content"])
        location_count = sum(1 for s in samples if "位置信息不一致" in s["messages"][1]["content"])
        hca_count = sum(1 for s in samples if "高后果区编号不一致" in s["messages"][1]["content"])
        
        print(f"样本分布：区段长不一致 {leader_count} 个，位置不一致 {location_count} 个，编号不一致 {hca_count} 个")
        print("="*80 + "\n")
    else:
        print(f"\n⚠️ 处理结束，无生成数据。请检查文件或API配置。")